package com.jiazixuan.stock.sharding;

import com.google.common.collect.Range;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingValue;
import org.joda.time.DateTime;

import java.util.Collection;
import java.util.Date;
import java.util.Optional;
import java.util.stream.Collectors;

/**
 * 定义个股流水表的分表算法类(覆盖个股)
 */
public class CommonAlg4Tb implements PreciseShardingAlgorithm<Date>, RangeShardingAlgorithm<Date> {
    /**
     * 分表策略：按照年月分表
     * 精准查询时走该方法：cur_time条件必须是 = 或者 in
     * eg：select * from stock_market_info where cur_time=xxxxx and code=xxxx
     * @param tbNames 所有可匹配表的集合 stock_rt_info_202101....stock_rt_info_202112
     *                                stock_rt_info_202201....stock_rt_info_202212
     * @param shardingValue
     * @return
     */
    @Override
    public String doSharding(Collection<String> tbNames, PreciseShardingValue<Date> shardingValue) {
        //获取逻辑表
        String logicTableName = shardingValue.getLogicTableName();
        //分片键 cur_time
        String columnName = shardingValue.getColumnName();
        //获取等值查询的条件值
        Date curTime = shardingValue.getValue();
        String yearMonth = new DateTime(curTime).toString("yyyyMM");
        Optional<String> result = tbNames.stream().filter(tbName -> tbName.endsWith(yearMonth)).findFirst();
        if (result.isPresent()) {
            return result.get();
        }
        return null;
    }

    /**
     * 范围查询
     * 分表策略：按照年月分表
     * eg：select * from stock_market_info where cur_time between xxxxx and xxx and code=xxxx
     * @param tbNames 所有可匹配表的集合 stock_rt_info_202101....stock_rt_info_202112
     *                                stock_rt_info_202201....stock_rt_info_202212
     * @param shardingValue
     * @return
     */
    @Override
    public Collection<String> doSharding(Collection<String> tbNames, RangeShardingValue<Date> shardingValue) {
        //获取逻辑表
        String logicTableName = shardingValue.getLogicTableName();
        //分片键 cur_time
        String columnName = shardingValue.getColumnName();
        //获取范围条件值
        Range<Date> valueRange = shardingValue.getValueRange();
        //判断下限
        if (valueRange.hasLowerBound()) {
            //获取起始值
            Date startTime = valueRange.lowerEndpoint();
            //获取条件值初始时间年份
            int startYearMonth = Integer.parseInt(new DateTime(startTime).toString("yyyyMM"));
            //过滤出条件值中大于起始值的年份
            tbNames = tbNames.stream().filter(tbName->Integer.parseInt(tbName.substring(tbName.lastIndexOf("_")+1))>=startYearMonth).collect(Collectors.toList());
        }
        //判断上限
        if (valueRange.hasUpperBound()) {
            //获取结束值
            Date endTime = valueRange.upperEndpoint();
            //获取条件值初始时间年份
            int endYearMonth = Integer.parseInt(new DateTime(endTime).toString("yyyyMM"));
            //过滤出条件值中小于结束值的年份
            tbNames = tbNames.stream().filter(tbName->Integer.parseInt(tbName.substring(tbName.lastIndexOf("_")+1))<=endYearMonth).collect(Collectors.toList());
        }
        return tbNames;
    }
}
